use bytes::{Buf, BufMut};
use serde::{de::DeserializeOwned, Serialize};
use std::{io::Write, marker::PhantomData};
use tokio_util::codec::{Decoder, Encoder};

#[derive(Debug)]
pub(crate) enum EncodeError<SE> {
    IOError(std::io::Error),
    SerializeError(SE),
}
#[derive(Debug)]
pub(crate) enum DecodeError<DE> {
    IOError(std::io::Error),
    DeserializeError(DE),
}

impl<SE> From<std::io::Error> for EncodeError<SE> {
    fn from(value: std::io::Error) -> Self {
        Self::IOError(value)
    }
}
impl<DE> From<std::io::Error> for DecodeError<DE> {
    fn from(value: std::io::Error) -> Self {
        Self::IOError(value)
    }
}

pub(crate) struct ChannelEndCodec<I, O, SD> {
    decode_payload_size: Option<usize>,
    serder: SD,
    marker: PhantomData<(fn(I) -> I, fn(O) -> O)>,
}

impl<I, O, SD> ChannelEndCodec<I, O, SD> {
    pub(crate) fn new(serder: SD) -> Self {
        Self {
            decode_payload_size: None,
            serder,
            marker: std::marker::PhantomData,
        }
    }
}

pub trait ToWriterSerializer {
    type Err;
    fn serialize_one_to_writer<W: Write, T: Serialize>(
        &mut self,
        w: W,
        t: T,
    ) -> Result<(), Self::Err>;
}

pub trait FromSliceDeserializer {
    type Err;
    fn deserialize_one_from_slice<T: DeserializeOwned>(&mut self, s: &[u8])
        -> Result<T, Self::Err>;
}

impl<I: Serialize, O, SD: ToWriterSerializer> Encoder<I> for ChannelEndCodec<I, O, SD> {
    type Error = EncodeError<SD::Err>;

    fn encode(&mut self, item: I, dst: &mut bytes::BytesMut) -> Result<(), Self::Error> {
        let orig_len = dst.len();
        dst.reserve(16);
        unsafe { dst.advance_mut(8) };
        let mut payload_writer = dst.writer();
        self.serder
            .serialize_one_to_writer::<_, I>(&mut payload_writer, item)
            .map_err(EncodeError::SerializeError)?;
        let len = dst.len() - orig_len - 8;
        dst[orig_len..orig_len + 8].copy_from_slice(&(len.to_le_bytes()));
        Ok(())
    }
}

impl<I, O, SD> Decoder for ChannelEndCodec<I, O, SD>
where
    O: DeserializeOwned,
    SD: FromSliceDeserializer,
{
    type Error = DecodeError<SD::Err>;
    type Item = O;

    fn decode(&mut self, src: &mut bytes::BytesMut) -> Result<Option<Self::Item>, Self::Error> {
        loop {
            if let Some(payload_sz) = self.decode_payload_size {
                if src.len() < payload_sz {
                    return Ok(None);
                } else {
                    self.decode_payload_size = None;
                    let out = self
                        .serder
                        .deserialize_one_from_slice::<O>(&src[..payload_sz])
                        .map_err(DecodeError::DeserializeError)?;
                    src.advance(payload_sz);
                    return Ok(Some(out));
                }
            } else if src.len() < 8 {
                return Ok(None);
            } else {
                let size = src.get_u64_le();
                self.decode_payload_size = Some(size as usize);
            }
        }
    }
}

#[cfg(feature = "msgpack")]
pub use self::msgpack::DefaultSerder;

#[cfg(feature = "msgpack")]
mod msgpack {
    use super::{FromSliceDeserializer, ToWriterSerializer};
    use serde::{de::DeserializeOwned, Serialize};
    use std::io::Write;
    #[derive(Default, Clone)]
    pub struct DefaultSerder;
    impl ToWriterSerializer for DefaultSerder {
        type Err = rmp_serde::encode::Error;

        fn serialize_one_to_writer<W: Write, T: Serialize>(
            &mut self,
            w: W,
            t: T,
        ) -> Result<(), Self::Err> {
            let mut ser = rmp_serde::Serializer::new(w);
            t.serialize(&mut ser)?;
            Ok(())
        }
    }
    impl FromSliceDeserializer for DefaultSerder {
        type Err = rmp_serde::decode::Error;

        fn deserialize_one_from_slice<T: DeserializeOwned>(
            &mut self,
            s: &[u8],
        ) -> Result<T, Self::Err> {
            rmp_serde::from_slice(s)
        }
    }
}
